/*
 * IKEv2 initiator (RFC 4306) for EAP-IKEV2
 * Copyright (c) 2007, Jouni Malinen <j@w1.fi>
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License version 2 as
 * published by the Free Software Foundation.
 *
 * Alternatively, this software may be distributed under the terms of BSD
 * license.
 *
 * See README and COPYING for more details.
 */

#include "includes.h"

#include "common.h"
#include "dh_groups.h"
#include "ikev2.h"


static int ikev2_process_idr(struct ikev2_initiator_data *data,
                             const u8 *idr, size_t idr_len);


void ikev2_initiator_deinit(struct ikev2_initiator_data *data)
{
        ikev2_free_keys(&data->keys);
        wpabuf_free(data->r_dh_public);
        wpabuf_free(data->i_dh_private);
        os_free(data->IDi);
        os_free(data->IDr);
        os_free(data->shared_secret);
        wpabuf_free(data->i_sign_msg);
        wpabuf_free(data->r_sign_msg);
        os_free(data->key_pad);
}


static int ikev2_derive_keys(struct ikev2_initiator_data *data)
{
        u8 *buf, *pos, *pad, skeyseed[IKEV2_MAX_HASH_LEN];
        size_t buf_len, pad_len;
        struct wpabuf *shared;
        const struct ikev2_integ_alg *integ;
        const struct ikev2_prf_alg *prf;
        const struct ikev2_encr_alg *encr;
        int ret;
        const u8 *addr[2];
        size_t len[2];

        /* RFC 4306, Sect. 2.14 */

        integ = ikev2_get_integ(data->proposal.integ);
        prf = ikev2_get_prf(data->proposal.prf);
        encr = ikev2_get_encr(data->proposal.encr);
        if (integ == NULL || prf == NULL || encr == NULL) {
                wpa_printf(MSG_INFO, "IKEV2: Unsupported proposal");
                return -1;
        }

        shared = dh_derive_shared(data->r_dh_public, data->i_dh_private,
                                  data->dh);
        if (shared == NULL)
                return -1;

        /* Construct Ni | Nr | SPIi | SPIr */

        buf_len = data->i_nonce_len + data->r_nonce_len + 2 * IKEV2_SPI_LEN;
        buf = os_malloc(buf_len);
        if (buf == NULL) {
                wpabuf_free(shared);
                return -1;
        }

        pos = buf;
        os_memcpy(pos, data->i_nonce, data->i_nonce_len);
        pos += data->i_nonce_len;
        os_memcpy(pos, data->r_nonce, data->r_nonce_len);
        pos += data->r_nonce_len;
        os_memcpy(pos, data->i_spi, IKEV2_SPI_LEN);
        pos += IKEV2_SPI_LEN;
        os_memcpy(pos, data->r_spi, IKEV2_SPI_LEN);

        /* SKEYSEED = prf(Ni | Nr, g^ir) */

        /* Use zero-padding per RFC 4306, Sect. 2.14 */
        pad_len = data->dh->prime_len - wpabuf_len(shared);
        pad = os_zalloc(pad_len ? pad_len : 1);
        if (pad == NULL) {
                wpabuf_free(shared);
                os_free(buf);
                return -1;
        }
        addr[0] = pad;
        len[0] = pad_len;
        addr[1] = wpabuf_head(shared);
        len[1] = wpabuf_len(shared);
        if (ikev2_prf_hash(prf->id, buf, data->i_nonce_len + data->r_nonce_len,
                           2, addr, len, skeyseed) < 0) {
                wpabuf_free(shared);
                os_free(buf);
                os_free(pad);
                return -1;
        }
        os_free(pad);
        wpabuf_free(shared);

        /* DH parameters are not needed anymore, so free them */
        wpabuf_free(data->r_dh_public);
        data->r_dh_public = NULL;
        wpabuf_free(data->i_dh_private);
        data->i_dh_private = NULL;

        wpa_hexdump_key(MSG_DEBUG, "IKEV2: SKEYSEED",
                        skeyseed, prf->hash_len);

        ret = ikev2_derive_sk_keys(prf, integ, encr, skeyseed, buf, buf_len,
                                   &data->keys);
        os_free(buf);
        return ret;
}


static int ikev2_parse_transform(struct ikev2_initiator_data *data,
                                 struct ikev2_proposal_data *prop,
                                 const u8 *pos, const u8 *end)
{
        int transform_len;
        const struct ikev2_transform *t;
        u16 transform_id;
        const u8 *tend;

        if (end - pos < (int) sizeof(*t)) {
                wpa_printf(MSG_INFO, "IKEV2: Too short transform");
                return -1;
        }

        t = (const struct ikev2_transform *) pos;
        transform_len = WPA_GET_BE16(t->transform_length);
        if (transform_len < (int) sizeof(*t) || pos + transform_len > end) {
                wpa_printf(MSG_INFO, "IKEV2: Invalid transform length %d",
                           transform_len);
                return -1;
        }
        tend = pos + transform_len;

        transform_id = WPA_GET_BE16(t->transform_id);

        wpa_printf(MSG_DEBUG, "IKEV2:   Transform:");
        wpa_printf(MSG_DEBUG, "IKEV2:     Type: %d  Transform Length: %d  "
                   "Transform Type: %d  Transform ID: %d",
                   t->type, transform_len, t->transform_type, transform_id);

        if (t->type != 0 && t->type != 3) {
                wpa_printf(MSG_INFO, "IKEV2: Unexpected Transform type");
                return -1;
        }

        pos = (const u8 *) (t + 1);
        if (pos < tend) {
                wpa_hexdump(MSG_DEBUG, "IKEV2:     Transform Attributes",
                            pos, tend - pos);
        }

        switch (t->transform_type) {
        case IKEV2_TRANSFORM_ENCR:
                if (ikev2_get_encr(transform_id) &&
                    transform_id == data->proposal.encr) {
                        if (transform_id == ENCR_AES_CBC) {
                                if (tend - pos != 4) {
                                        wpa_printf(MSG_DEBUG, "IKEV2: No "
                                                   "Transform Attr for AES");
                                        break;
                                }
                                if (WPA_GET_BE16(pos) != 0x800e) {
                                        wpa_printf(MSG_DEBUG, "IKEV2: Not a "
                                                   "Key Size attribute for "
                                                   "AES");
                                        break;
                                }
                                if (WPA_GET_BE16(pos + 2) != 128) {
                                        wpa_printf(MSG_DEBUG, "IKEV2: "
                                                   "Unsupported AES key size "
                                                   "%d bits",
                                                   WPA_GET_BE16(pos + 2));
                                        break;
                                }
                        }
                        prop->encr = transform_id;
                }
                break;
        case IKEV2_TRANSFORM_PRF:
                if (ikev2_get_prf(transform_id) &&
                    transform_id == data->proposal.prf)
                        prop->prf = transform_id;
                break;
        case IKEV2_TRANSFORM_INTEG:
                if (ikev2_get_integ(transform_id) &&
                    transform_id == data->proposal.integ)
                        prop->integ = transform_id;
                break;
        case IKEV2_TRANSFORM_DH:
                if (dh_groups_get(transform_id) &&
                    transform_id == data->proposal.dh)
                        prop->dh = transform_id;
                break;
        }

        return transform_len;
}


static int ikev2_parse_proposal(struct ikev2_initiator_data *data,
                                struct ikev2_proposal_data *prop,
                                const u8 *pos, const u8 *end)
{
        const u8 *pend, *ppos;
        int proposal_len, i;
        const struct ikev2_proposal *p;

        if (end - pos < (int) sizeof(*p)) {
                wpa_printf(MSG_INFO, "IKEV2: Too short proposal");
                return -1;
        }

        p = (const struct ikev2_proposal *) pos;
        proposal_len = WPA_GET_BE16(p->proposal_length);
        if (proposal_len < (int) sizeof(*p) || pos + proposal_len > end) {
                wpa_printf(MSG_INFO, "IKEV2: Invalid proposal length %d",
                           proposal_len);
                return -1;
        }
        wpa_printf(MSG_DEBUG, "IKEV2: SAi1 Proposal # %d",
                   p->proposal_num);
        wpa_printf(MSG_DEBUG, "IKEV2:   Type: %d  Proposal Length: %d "
                   " Protocol ID: %d",
                   p->type, proposal_len, p->protocol_id);
        wpa_printf(MSG_DEBUG, "IKEV2:   SPI Size: %d  Transforms: %d",
                   p->spi_size, p->num_transforms);

        if (p->type != 0 && p->type != 2) {
                wpa_printf(MSG_INFO, "IKEV2: Unexpected Proposal type");
                return -1;
        }

        if (p->protocol_id != IKEV2_PROTOCOL_IKE) {
                wpa_printf(MSG_DEBUG, "IKEV2: Unexpected Protocol ID "
                           "(only IKE allowed for EAP-IKEv2)");
                return -1;
        }

        if (p->proposal_num != prop->proposal_num) {
                if (p->proposal_num == prop->proposal_num + 1)
                        prop->proposal_num = p->proposal_num;
                else {
                        wpa_printf(MSG_INFO, "IKEV2: Unexpected Proposal #");
                        return -1;
                }
        }

        ppos = (const u8 *) (p + 1);
        pend = pos + proposal_len;
        if (ppos + p->spi_size > pend) {
                wpa_printf(MSG_INFO, "IKEV2: Not enough room for SPI "
                           "in proposal");
                return -1;
        }
        if (p->spi_size) {
                wpa_hexdump(MSG_DEBUG, "IKEV2:    SPI",
                            ppos, p->spi_size);
                ppos += p->spi_size;
        }

        /*
         * For initial IKE_SA negotiation, SPI Size MUST be zero; for
         * subsequent negotiations, it must be 8 for IKE. We only support
         * initial case for now.
         */
        if (p->spi_size != 0) {
                wpa_printf(MSG_INFO, "IKEV2: Unexpected SPI Size");
                return -1;
        }

        if (p->num_transforms == 0) {
                wpa_printf(MSG_INFO, "IKEV2: At least one transform required");
                return -1;
        }

        for (i = 0; i < (int) p->num_transforms; i++) {
                int tlen = ikev2_parse_transform(data, prop, ppos, pend);
                if (tlen < 0)
                        return -1;
                ppos += tlen;
        }

        if (ppos != pend) {
                wpa_printf(MSG_INFO, "IKEV2: Unexpected data after "
                           "transforms");
                return -1;
        }

        return proposal_len;
}


static int ikev2_process_sar1(struct ikev2_initiator_data *data,
                              const u8 *sar1, size_t sar1_len)
{
        struct ikev2_proposal_data prop;
        const u8 *pos, *end;
        int found = 0;

        /* Security Association Payloads: <Proposals> */

        if (sar1 == NULL) {
                wpa_printf(MSG_INFO, "IKEV2: SAr1 not received");
                return -1;
        }

        os_memset(&prop, 0, sizeof(prop));
        prop.proposal_num = 1;

        pos = sar1;
        end = sar1 + sar1_len;

        while (pos < end) {
                int plen;

                prop.integ = -1;
                prop.prf = -1;
                prop.encr = -1;
                prop.dh = -1;
                plen = ikev2_parse_proposal(data, &prop, pos, end);
                if (plen < 0)
                        return -1;

                if (!found && prop.integ != -1 && prop.prf != -1 &&
                    prop.encr != -1 && prop.dh != -1) {
                        found = 1;
                }

                pos += plen;

                /* Only one proposal expected in SAr */
                break;
        }

        if (pos != end) {
                wpa_printf(MSG_INFO, "IKEV2: Unexpected data after proposal");
                return -1;
        }

        if (!found) {
                wpa_printf(MSG_INFO, "IKEV2: No acceptable proposal found");
                return -1;
        }

        wpa_printf(MSG_DEBUG, "IKEV2: Accepted proposal #%d: ENCR:%d PRF:%d "
                   "INTEG:%d D-H:%d", data->proposal.proposal_num,
                   data->proposal.encr, data->proposal.prf,
                   data->proposal.integ, data->proposal.dh);

        return 0;
}


static int ikev2_process_ker(struct ikev2_initiator_data *data,
                             const u8 *ker, size_t ker_len)
{
        u16 group;

        /*
         * Key Exchange Payload:
         * DH Group # (16 bits)
         * RESERVED (16 bits)
         * Key Exchange Data (Diffie-Hellman public value)
         */

        if (ker == NULL) {
                wpa_printf(MSG_INFO, "IKEV2: KEr not received");
                return -1;
        }

        if (ker_len < 4 + 96) {
                wpa_printf(MSG_INFO, "IKEV2: Too show Key Exchange Payload");
                return -1;
        }

        group = WPA_GET_BE16(ker);
        wpa_printf(MSG_DEBUG, "IKEV2: KEr DH Group #%u", group);

        if (group != data->proposal.dh) {
                wpa_printf(MSG_DEBUG, "IKEV2: KEr DH Group #%u does not match "
                           "with the selected proposal (%u)",
                           group, data->proposal.dh);
                return -1;
        }

        if (data->dh == NULL) {
                wpa_printf(MSG_INFO, "IKEV2: Unsupported DH group");
                return -1;
        }

        /* RFC 4306, Section 3.4:
         * The length of DH public value MUST be equal to the lenght of the
         * prime modulus.
         */
        if (ker_len - 4 != data->dh->prime_len) {
                wpa_printf(MSG_INFO, "IKEV2: Invalid DH public value length "
                           "%ld (expected %ld)",
                           (long) (ker_len - 4), (long) data->dh->prime_len);
                return -1;
        }

        wpabuf_free(data->r_dh_public);
        data->r_dh_public = wpabuf_alloc_copy(ker + 4, ker_len - 4);
        if (data->r_dh_public == NULL)
                return -1;

        wpa_hexdump_buf(MSG_DEBUG, "IKEV2: KEr Diffie-Hellman Public Value",
                        data->r_dh_public);
        
        return 0;
}


static int ikev2_process_nr(struct ikev2_initiator_data *data,
                            const u8 *nr, size_t nr_len)
{
        if (nr == NULL) {
                wpa_printf(MSG_INFO, "IKEV2: Nr not received");
                return -1;
        }

        if (nr_len < IKEV2_NONCE_MIN_LEN || nr_len > IKEV2_NONCE_MAX_LEN) {
                wpa_printf(MSG_INFO, "IKEV2: Invalid Nr length %ld",
                           (long) nr_len);
                return -1;
        }

        data->r_nonce_len = nr_len;
        os_memcpy(data->r_nonce, nr, nr_len);
        wpa_hexdump(MSG_MSGDUMP, "IKEV2: Nr",
                    data->r_nonce, data->r_nonce_len);

        return 0;
}


static int ikev2_process_sa_init_encr(struct ikev2_initiator_data *data,
                                      const struct ikev2_hdr *hdr,
                                      const u8 *encrypted,
                                      size_t encrypted_len, u8 next_payload)
{
        u8 *decrypted;
        size_t decrypted_len;
        struct ikev2_payloads pl;
        int ret = 0;

        decrypted = ikev2_decrypt_payload(data->proposal.encr,
                                          data->proposal.integ, &data->keys, 0,
                                          hdr, encrypted, encrypted_len,
                                          &decrypted_len);
        if (decrypted == NULL)
                return -1;

        wpa_printf(MSG_DEBUG, "IKEV2: Processing decrypted payloads");

        if (ikev2_parse_payloads(&pl, next_payload, decrypted,
                                 decrypted + decrypted_len) < 0) {
                wpa_printf(MSG_INFO, "IKEV2: Failed to parse decrypted "
                           "payloads");
                return -1;
        }

        if (pl.idr)
                ret = ikev2_process_idr(data, pl.idr, pl.idr_len);

        os_free(decrypted);

        return ret;
}


static int ikev2_process_sa_init(struct ikev2_initiator_data *data,
                                 const struct ikev2_hdr *hdr,
                                 struct ikev2_payloads *pl)
{
        if (ikev2_process_sar1(data, pl->sa, pl->sa_len) < 0 ||
            ikev2_process_ker(data, pl->ke, pl->ke_len) < 0 ||
            ikev2_process_nr(data, pl->nonce, pl->nonce_len) < 0)
                return -1;

        os_memcpy(data->r_spi, hdr->r_spi, IKEV2_SPI_LEN);

        if (ikev2_derive_keys(data) < 0)
                return -1;

        if (pl->encrypted) {
                wpa_printf(MSG_DEBUG, "IKEV2: Encrypted payload in SA_INIT - "
                           "try to get IDr from it");
                if (ikev2_process_sa_init_encr(data, hdr, pl->encrypted,
                                               pl->encrypted_len,
                                               pl->encr_next_payload) < 0) {
                        wpa_printf(MSG_INFO, "IKEV2: Failed to process "
                                   "encrypted payload");
                        return -1;
                }
        }

        data->state = SA_AUTH;

        return 0;
}


static int ikev2_process_idr(struct ikev2_initiator_data *data,
                             const u8 *idr, size_t idr_len)
{
        u8 id_type;

        if (idr == NULL) {
                wpa_printf(MSG_INFO, "IKEV2: No IDr received");
                return -1;
        }

        if (idr_len < 4) {
                wpa_printf(MSG_INFO, "IKEV2: Too short IDr payload");
                return -1;
        }

        id_type = idr[0];
        idr += 4;
        idr_len -= 4;

        wpa_printf(MSG_DEBUG, "IKEV2: IDr ID Type %d", id_type);
        wpa_hexdump_ascii(MSG_DEBUG, "IKEV2: IDr", idr, idr_len);
        if (data->IDr) {
                if (id_type != data->IDr_type || idr_len != data->IDr_len ||
                    os_memcmp(idr, data->IDr, idr_len) != 0) {
                        wpa_printf(MSG_INFO, "IKEV2: IDr differs from the one "
                                   "received earlier");
                        wpa_printf(MSG_DEBUG, "IKEV2: Previous IDr ID Type %d",
                                   id_type);
                        wpa_hexdump_ascii(MSG_DEBUG, "Previous IKEV2: IDr",
                                          data->IDr, data->IDr_len);
                        return -1;
                }
                os_free(data->IDr);
        }
        data->IDr = os_malloc(idr_len);
        if (data->IDr == NULL)
                return -1;
        os_memcpy(data->IDr, idr, idr_len);
        data->IDr_len = idr_len;
        data->IDr_type = id_type;

        return 0;
}


static int ikev2_process_cert(struct ikev2_initiator_data *data,
                              const u8 *cert, size_t cert_len)
{
        u8 cert_encoding;

        if (cert == NULL) {
                if (data->peer_auth == PEER_AUTH_CERT) {
                        wpa_printf(MSG_INFO, "IKEV2: No Certificate received");
                        return -1;
                }
                return 0;
        }

        if (cert_len < 1) {
                wpa_printf(MSG_INFO, "IKEV2: No Cert Encoding field");
                return -1;
        }

        cert_encoding = cert[0];
        cert++;
        cert_len--;

        wpa_printf(MSG_DEBUG, "IKEV2: Cert Encoding %d", cert_encoding);
        wpa_hexdump(MSG_MSGDUMP, "IKEV2: Certificate Data", cert, cert_len);

        /* TODO: validate certificate */

        return 0;
}


static int ikev2_process_auth_cert(struct ikev2_initiator_data *data,
                                   u8 method, const u8 *auth, size_t auth_len)
{
        if (method != AUTH_RSA_SIGN) {
                wpa_printf(MSG_INFO, "IKEV2: Unsupported authentication "
                           "method %d", method);
                return -1;
        }

        /* TODO: validate AUTH */
        return 0;
}


static int ikev2_process_auth_secret(struct ikev2_initiator_data *data,
                                     u8 method, const u8 *auth,
                                     size_t auth_len)
{
        u8 auth_data[IKEV2_MAX_HASH_LEN];
        const struct ikev2_prf_alg *prf;

        if (method != AUTH_SHARED_KEY_MIC) {
                wpa_printf(MSG_INFO, "IKEV2: Unsupported authentication "
                           "method %d", method);
                return -1;
        }

        /* msg | Ni | prf(SK_pr,IDr') */
        if (ikev2_derive_auth_data(data->proposal.prf, data->r_sign_msg,
                                   data->IDr, data->IDr_len, data->IDr_type,
                                   &data->keys, 0, data->shared_secret,
                                   data->shared_secret_len,
                                   data->i_nonce, data->i_nonce_len,
                                   data->key_pad, data->key_pad_len,
                                   auth_data) < 0) {
                wpa_printf(MSG_INFO, "IKEV2: Could not derive AUTH data");
                return -1;
        }

        wpabuf_free(data->r_sign_msg);
        data->r_sign_msg = NULL;

        prf = ikev2_get_prf(data->proposal.prf);
        if (prf == NULL)
                return -1;

        if (auth_len != prf->hash_len ||
            os_memcmp(auth, auth_data, auth_len) != 0) {
                wpa_printf(MSG_INFO, "IKEV2: Invalid Authentication Data");
                wpa_hexdump(MSG_DEBUG, "IKEV2: Received Authentication Data",
                            auth, auth_len);
                wpa_hexdump(MSG_DEBUG, "IKEV2: Expected Authentication Data",
                            auth_data, prf->hash_len);
                return -1;
        }

        wpa_printf(MSG_DEBUG, "IKEV2: Peer authenticated successfully "
                   "using shared keys");

        return 0;
}


static int ikev2_process_auth(struct ikev2_initiator_data *data,
                              const u8 *auth, size_t auth_len)
{
        u8 auth_method;

        if (auth == NULL) {
                wpa_printf(MSG_INFO, "IKEV2: No Authentication Payload");
                return -1;
        }

        if (auth_len < 4) {
                wpa_printf(MSG_INFO, "IKEV2: Too short Authentication "
                           "Payload");
                return -1;
        }

        auth_method = auth[0];
        auth += 4;
        auth_len -= 4;

        wpa_printf(MSG_DEBUG, "IKEV2: Auth Method %d", auth_method);
        wpa_hexdump(MSG_MSGDUMP, "IKEV2: Authentication Data", auth, auth_len);

        switch (data->peer_auth) {
        case PEER_AUTH_CERT:
                return ikev2_process_auth_cert(data, auth_method, auth,
                                               auth_len);
        case PEER_AUTH_SECRET:
                return ikev2_process_auth_secret(data, auth_method, auth,
                                                 auth_len);
        }

        return -1;
}


static int ikev2_process_sa_auth_decrypted(struct ikev2_initiator_data *data,
                                           u8 next_payload,
                                           u8 *payload, size_t payload_len)
{
        struct ikev2_payloads pl;

        wpa_printf(MSG_DEBUG, "IKEV2: Processing decrypted payloads");

        if (ikev2_parse_payloads(&pl, next_payload, payload, payload +
                                 payload_len) < 0) {
                wpa_printf(MSG_INFO, "IKEV2: Failed to parse decrypted "
                           "payloads");
                return -1;
        }

        if (ikev2_process_idr(data, pl.idr, pl.idr_len) < 0 ||
            ikev2_process_cert(data, pl.cert, pl.cert_len) < 0 ||
            ikev2_process_auth(data, pl.auth, pl.auth_len) < 0)
                return -1;

        return 0;
}


static int ikev2_process_sa_auth(struct ikev2_initiator_data *data,
                                 const struct ikev2_hdr *hdr,
                                 struct ikev2_payloads *pl)
{
        u8 *decrypted;
        size_t decrypted_len;
        int ret;

        decrypted = ikev2_decrypt_payload(data->proposal.encr,
                                          data->proposal.integ,
                                          &data->keys, 0, hdr, pl->encrypted,
                                          pl->encrypted_len, &decrypted_len);
        if (decrypted == NULL)
                return -1;

        ret = ikev2_process_sa_auth_decrypted(data, pl->encr_next_payload,
                                              decrypted, decrypted_len);
        os_free(decrypted);

        if (ret == 0 && !data->unknown_user) {
                wpa_printf(MSG_DEBUG, "IKEV2: Authentication completed");
                data->state = IKEV2_DONE;
        }

        return ret;
}


static int ikev2_validate_rx_state(struct ikev2_initiator_data *data,
                                   u8 exchange_type, u32 message_id)
{
        switch (data->state) {
        case SA_INIT:
                /* Expect to receive IKE_SA_INIT: HDR, SAr, KEr, Nr, [CERTREQ],
                 * [SK{IDr}] */
                if (exchange_type != IKE_SA_INIT) {
                        wpa_printf(MSG_INFO, "IKEV2: Unexpected Exchange Type "
                                   "%u in SA_INIT state", exchange_type);
                        return -1;
                }
                if (message_id != 0) {
                        wpa_printf(MSG_INFO, "IKEV2: Unexpected Message ID %u "
                                   "in SA_INIT state", message_id);
                        return -1;
                }
                break;
        case SA_AUTH:
                /* Expect to receive IKE_SA_AUTH:
                 * HDR, SK {IDr, [CERT,] [CERTREQ,] [NFID,] AUTH}
                 */
                if (exchange_type != IKE_SA_AUTH) {
                        wpa_printf(MSG_INFO, "IKEV2: Unexpected Exchange Type "
                                   "%u in SA_AUTH state", exchange_type);
                        return -1;
                }
                if (message_id != 1) {
                        wpa_printf(MSG_INFO, "IKEV2: Unexpected Message ID %u "
                                   "in SA_AUTH state", message_id);
                        return -1;
                }
                break;
        case CHILD_SA:
                if (exchange_type != CREATE_CHILD_SA) {
                        wpa_printf(MSG_INFO, "IKEV2: Unexpected Exchange Type "
                                   "%u in CHILD_SA state", exchange_type);
                        return -1;
                }
                if (message_id != 2) {
                        wpa_printf(MSG_INFO, "IKEV2: Unexpected Message ID %u "
                                   "in CHILD_SA state", message_id);
                        return -1;
                }
                break;
        case IKEV2_DONE:
                return -1;
        }

        return 0;
}


int ikev2_initiator_process(struct ikev2_initiator_data *data,
                            const struct wpabuf *buf)
{
        const struct ikev2_hdr *hdr;
        u32 length, message_id;
        const u8 *pos, *end;
        struct ikev2_payloads pl;

        wpa_printf(MSG_MSGDUMP, "IKEV2: Received message (len %lu)",
                   (unsigned long) wpabuf_len(buf));

        if (wpabuf_len(buf) < sizeof(*hdr)) {
                wpa_printf(MSG_INFO, "IKEV2: Too short frame to include HDR");
                return -1;
        }

        hdr = (const struct ikev2_hdr *) wpabuf_head(buf);
        end = wpabuf_head_u8(buf) + wpabuf_len(buf);
        message_id = WPA_GET_BE32(hdr->message_id);
        length = WPA_GET_BE32(hdr->length);

        wpa_hexdump(MSG_DEBUG, "IKEV2:   IKE_SA Initiator's SPI",
                    hdr->i_spi, IKEV2_SPI_LEN);
        wpa_hexdump(MSG_DEBUG, "IKEV2:   IKE_SA Initiator's SPI",
                    hdr->r_spi, IKEV2_SPI_LEN);
        wpa_printf(MSG_DEBUG, "IKEV2:   Next Payload: %u  Version: 0x%x  "
                   "Exchange Type: %u",
                   hdr->next_payload, hdr->version, hdr->exchange_type);
        wpa_printf(MSG_DEBUG, "IKEV2:   Message ID: %u  Length: %u",
                   message_id, length);

        if (hdr->version != IKEV2_VERSION) {
                wpa_printf(MSG_INFO, "IKEV2: Unsupported HDR version 0x%x "
                           "(expected 0x%x)", hdr->version, IKEV2_VERSION);
                return -1;
        }

        if (length != wpabuf_len(buf)) {
                wpa_printf(MSG_INFO, "IKEV2: Invalid length (HDR: %lu != "
                           "RX: %lu)", (unsigned long) length,
                           (unsigned long) wpabuf_len(buf));
                return -1;
        }

        if (ikev2_validate_rx_state(data, hdr->exchange_type, message_id) < 0)
                return -1;

        if ((hdr->flags & (IKEV2_HDR_INITIATOR | IKEV2_HDR_RESPONSE)) !=
            IKEV2_HDR_RESPONSE) {
                wpa_printf(MSG_INFO, "IKEV2: Unexpected Flags value 0x%x",
                           hdr->flags);
                return -1;
        }

        if (data->state != SA_INIT) {
                if (os_memcmp(data->i_spi, hdr->i_spi, IKEV2_SPI_LEN) != 0) {
                        wpa_printf(MSG_INFO, "IKEV2: Unexpected IKE_SA "
                                   "Initiator's SPI");
                        return -1;
                }
                if (os_memcmp(data->r_spi, hdr->r_spi, IKEV2_SPI_LEN) != 0) {
                        wpa_printf(MSG_INFO, "IKEV2: Unexpected IKE_SA "
                                   "Responder's SPI");
                        return -1;
                }
        }

        pos = (const u8 *) (hdr + 1);
        if (ikev2_parse_payloads(&pl, hdr->next_payload, pos, end) < 0)
                return -1;

        switch (data->state) {
        case SA_INIT:
                if (ikev2_process_sa_init(data, hdr, &pl) < 0)
                        return -1;
                wpabuf_free(data->r_sign_msg);
                data->r_sign_msg = wpabuf_dup(buf);
                break;
        case SA_AUTH:
                if (ikev2_process_sa_auth(data, hdr, &pl) < 0)
                        return -1;
                break;
        case CHILD_SA:
        case IKEV2_DONE:
                break;
        }

        return 0;
}


static void ikev2_build_hdr(struct ikev2_initiator_data *data,
                            struct wpabuf *msg, u8 exchange_type,
                            u8 next_payload, u32 message_id)
{
        struct ikev2_hdr *hdr;

        wpa_printf(MSG_DEBUG, "IKEV2: Adding HDR");

        /* HDR - RFC 4306, Sect. 3.1 */
        hdr = wpabuf_put(msg, sizeof(*hdr));
        os_memcpy(hdr->i_spi, data->i_spi, IKEV2_SPI_LEN);
        os_memcpy(hdr->r_spi, data->r_spi, IKEV2_SPI_LEN);
        hdr->next_payload = next_payload;
        hdr->version = IKEV2_VERSION;
        hdr->exchange_type = exchange_type;
        hdr->flags = IKEV2_HDR_INITIATOR;
        WPA_PUT_BE32(hdr->message_id, message_id);
}


static int ikev2_build_sai(struct ikev2_initiator_data *data,
                            struct wpabuf *msg, u8 next_payload)
{
        struct ikev2_payload_hdr *phdr;
        size_t plen;
        struct ikev2_proposal *p;
        struct ikev2_transform *t;

        wpa_printf(MSG_DEBUG, "IKEV2: Adding SAi payload");

        /* SAi1 - RFC 4306, Sect. 2.7 and 3.3 */
        phdr = wpabuf_put(msg, sizeof(*phdr));
        phdr->next_payload = next_payload;
        phdr->flags = 0;

        /* TODO: support for multiple proposals */
        p = wpabuf_put(msg, sizeof(*p));
        p->proposal_num = data->proposal.proposal_num;
        p->protocol_id = IKEV2_PROTOCOL_IKE;
        p->num_transforms = 4;

        t = wpabuf_put(msg, sizeof(*t));
        t->type = 3;
        t->transform_type = IKEV2_TRANSFORM_ENCR;
        WPA_PUT_BE16(t->transform_id, data->proposal.encr);
        if (data->proposal.encr == ENCR_AES_CBC) {
                /* Transform Attribute: Key Len = 128 bits */
                wpabuf_put_be16(msg, 0x800e); /* AF=1, AttrType=14 */
                wpabuf_put_be16(msg, 128); /* 128-bit key */
        }
        plen = (u8 *) wpabuf_put(msg, 0) - (u8 *) t;
        WPA_PUT_BE16(t->transform_length, plen);

        t = wpabuf_put(msg, sizeof(*t));
        t->type = 3;
        WPA_PUT_BE16(t->transform_length, sizeof(*t));
        t->transform_type = IKEV2_TRANSFORM_PRF;
        WPA_PUT_BE16(t->transform_id, data->proposal.prf);

        t = wpabuf_put(msg, sizeof(*t));
        t->type = 3;
        WPA_PUT_BE16(t->transform_length, sizeof(*t));
        t->transform_type = IKEV2_TRANSFORM_INTEG;
        WPA_PUT_BE16(t->transform_id, data->proposal.integ);

        t = wpabuf_put(msg, sizeof(*t));
        WPA_PUT_BE16(t->transform_length, sizeof(*t));
        t->transform_type = IKEV2_TRANSFORM_DH;
        WPA_PUT_BE16(t->transform_id, data->proposal.dh);

        plen = (u8 *) wpabuf_put(msg, 0) - (u8 *) p;
        WPA_PUT_BE16(p->proposal_length, plen);

        plen = (u8 *) wpabuf_put(msg, 0) - (u8 *) phdr;
        WPA_PUT_BE16(phdr->payload_length, plen);

        return 0;
}


static int ikev2_build_kei(struct ikev2_initiator_data *data,
                           struct wpabuf *msg, u8 next_payload)
{
        struct ikev2_payload_hdr *phdr;
        size_t plen;
        struct wpabuf *pv;

        wpa_printf(MSG_DEBUG, "IKEV2: Adding KEi payload");

        data->dh = dh_groups_get(data->proposal.dh);
        pv = dh_init(data->dh, &data->i_dh_private);
        if (pv == NULL) {
                wpa_printf(MSG_DEBUG, "IKEV2: Failed to initialize DH");
                return -1;
        }

        /* KEi - RFC 4306, Sect. 3.4 */
        phdr = wpabuf_put(msg, sizeof(*phdr));
        phdr->next_payload = next_payload;
        phdr->flags = 0;

        wpabuf_put_be16(msg, data->proposal.dh); /* DH Group # */
        wpabuf_put(msg, 2); /* RESERVED */
        /*
         * RFC 4306, Sect. 3.4: possible zero padding for public value to
         * match the length of the prime.
         */
        wpabuf_put(msg, data->dh->prime_len - wpabuf_len(pv));
        wpabuf_put_buf(msg, pv);
        os_free(pv);

        plen = (u8 *) wpabuf_put(msg, 0) - (u8 *) phdr;
        WPA_PUT_BE16(phdr->payload_length, plen);
        return 0;
}


static int ikev2_build_ni(struct ikev2_initiator_data *data,
                          struct wpabuf *msg, u8 next_payload)
{
        struct ikev2_payload_hdr *phdr;
        size_t plen;

        wpa_printf(MSG_DEBUG, "IKEV2: Adding Ni payload");

        /* Ni - RFC 4306, Sect. 3.9 */
        phdr = wpabuf_put(msg, sizeof(*phdr));
        phdr->next_payload = next_payload;
        phdr->flags = 0;
        wpabuf_put_data(msg, data->i_nonce, data->i_nonce_len);
        plen = (u8 *) wpabuf_put(msg, 0) - (u8 *) phdr;
        WPA_PUT_BE16(phdr->payload_length, plen);
        return 0;
}


static int ikev2_build_idi(struct ikev2_initiator_data *data,
                           struct wpabuf *msg, u8 next_payload)
{
        struct ikev2_payload_hdr *phdr;
        size_t plen;

        wpa_printf(MSG_DEBUG, "IKEV2: Adding IDi payload");

        if (data->IDi == NULL) {
                wpa_printf(MSG_INFO, "IKEV2: No IDi available");
                return -1;
        }

        /* IDi - RFC 4306, Sect. 3.5 */
        phdr = wpabuf_put(msg, sizeof(*phdr));
        phdr->next_payload = next_payload;
        phdr->flags = 0;
        wpabuf_put_u8(msg, ID_KEY_ID);
        wpabuf_put(msg, 3); /* RESERVED */
        wpabuf_put_data(msg, data->IDi, data->IDi_len);
        plen = (u8 *) wpabuf_put(msg, 0) - (u8 *) phdr;
        WPA_PUT_BE16(phdr->payload_length, plen);
        return 0;
}


static int ikev2_build_auth(struct ikev2_initiator_data *data,
                            struct wpabuf *msg, u8 next_payload)
{
        struct ikev2_payload_hdr *phdr;
        size_t plen;
        const struct ikev2_prf_alg *prf;

        wpa_printf(MSG_DEBUG, "IKEV2: Adding AUTH payload");

        prf = ikev2_get_prf(data->proposal.prf);
        if (prf == NULL)
                return -1;

        /* Authentication - RFC 4306, Sect. 3.8 */
        phdr = wpabuf_put(msg, sizeof(*phdr));
        phdr->next_payload = next_payload;
        phdr->flags = 0;
        wpabuf_put_u8(msg, AUTH_SHARED_KEY_MIC);
        wpabuf_put(msg, 3); /* RESERVED */

        /* msg | Nr | prf(SK_pi,IDi') */
        if (ikev2_derive_auth_data(data->proposal.prf, data->i_sign_msg,
                                   data->IDi, data->IDi_len, ID_KEY_ID,
                                   &data->keys, 1, data->shared_secret,
                                   data->shared_secret_len,
                                   data->r_nonce, data->r_nonce_len,
                                   data->key_pad, data->key_pad_len,
                                   wpabuf_put(msg, prf->hash_len)) < 0) {
                wpa_printf(MSG_INFO, "IKEV2: Could not derive AUTH data");
                return -1;
        }
        wpabuf_free(data->i_sign_msg);
        data->i_sign_msg = NULL;

        plen = (u8 *) wpabuf_put(msg, 0) - (u8 *) phdr;
        WPA_PUT_BE16(phdr->payload_length, plen);
        return 0;
}


static struct wpabuf * ikev2_build_sa_init(struct ikev2_initiator_data *data)
{
        struct wpabuf *msg;

        /* build IKE_SA_INIT: HDR, SAi, KEi, Ni */

        if (os_get_random(data->i_spi, IKEV2_SPI_LEN))
                return NULL;
        wpa_hexdump(MSG_DEBUG, "IKEV2: IKE_SA Initiator's SPI",
                    data->i_spi, IKEV2_SPI_LEN);

        data->i_nonce_len = IKEV2_NONCE_MIN_LEN;
        if (os_get_random(data->i_nonce, data->i_nonce_len))
                return NULL;
        wpa_hexdump(MSG_DEBUG, "IKEV2: Ni", data->i_nonce, data->i_nonce_len);

        msg = wpabuf_alloc(sizeof(struct ikev2_hdr) + 1000);
        if (msg == NULL)
                return NULL;

        ikev2_build_hdr(data, msg, IKE_SA_INIT, IKEV2_PAYLOAD_SA, 0);
        if (ikev2_build_sai(data, msg, IKEV2_PAYLOAD_KEY_EXCHANGE) ||
            ikev2_build_kei(data, msg, IKEV2_PAYLOAD_NONCE) ||
            ikev2_build_ni(data, msg, IKEV2_PAYLOAD_NO_NEXT_PAYLOAD)) {
                wpabuf_free(msg);
                return NULL;
        }

        ikev2_update_hdr(msg);

        wpa_hexdump_buf(MSG_MSGDUMP, "IKEV2: Sending message (SA_INIT)", msg);

        wpabuf_free(data->i_sign_msg);
        data->i_sign_msg = wpabuf_dup(msg);

        return msg;
}


static struct wpabuf * ikev2_build_sa_auth(struct ikev2_initiator_data *data)
{
        struct wpabuf *msg, *plain;
        const u8 *secret;
        size_t secret_len;

        secret = data->get_shared_secret(data->cb_ctx, data->IDr,
                                         data->IDr_len, &secret_len);
        if (secret == NULL) {
                wpa_printf(MSG_INFO, "IKEV2: Could not get shared secret - "
                           "use fake value");
                /* RFC 5106, Sect. 7:
                 * Use a random key to fake AUTH generation in order to prevent
                 * probing of user identities.
                 */
                data->unknown_user = 1;
                os_free(data->shared_secret);
                data->shared_secret = os_malloc(16);
                if (data->shared_secret == NULL)
                        return NULL;
                data->shared_secret_len = 16;
                if (os_get_random(data->shared_secret, 16))
                        return NULL;
        } else {
                os_free(data->shared_secret);
                data->shared_secret = os_malloc(secret_len);
                if (data->shared_secret == NULL)
                        return NULL;
                os_memcpy(data->shared_secret, secret, secret_len);
                data->shared_secret_len = secret_len;
        }

        /* build IKE_SA_AUTH: HDR, SK {IDi, [CERT,] [CERTREQ,] AUTH} */

        msg = wpabuf_alloc(sizeof(struct ikev2_hdr) + data->IDr_len + 1000);
        if (msg == NULL)
                return NULL;
        ikev2_build_hdr(data, msg, IKE_SA_AUTH, IKEV2_PAYLOAD_ENCRYPTED, 1);

        plain = wpabuf_alloc(data->IDr_len + 1000);
        if (plain == NULL) {
                wpabuf_free(msg);
                return NULL;
        }

        if (ikev2_build_idi(data, plain, IKEV2_PAYLOAD_AUTHENTICATION) ||
            ikev2_build_auth(data, plain, IKEV2_PAYLOAD_NO_NEXT_PAYLOAD) ||
            ikev2_build_encrypted(data->proposal.encr, data->proposal.integ,
                                  &data->keys, 1, msg, plain,
                                  IKEV2_PAYLOAD_IDi)) {
                wpabuf_free(plain);
                wpabuf_free(msg);
                return NULL;
        }
        wpabuf_free(plain);

        wpa_hexdump_buf(MSG_MSGDUMP, "IKEV2: Sending message (SA_AUTH)", msg);

        return msg;
}


struct wpabuf * ikev2_initiator_build(struct ikev2_initiator_data *data)
{
        switch (data->state) {
        case SA_INIT:
                return ikev2_build_sa_init(data);
        case SA_AUTH:
                return ikev2_build_sa_auth(data);
        case CHILD_SA:
                return NULL;
        case IKEV2_DONE:
                return NULL;
        }
        return NULL;
}
